package middleware

import (
	"bytes"
	"crypto/aes"
	"fmt"
	"io"
	"net/http"
	"runtime"

	"github.com/duke-git/lancet/cryptor"
	"github.com/gin-gonic/gin"
	"github.com/rs/zerolog/log"
	"github.com/valyala/bytebufferpool"
)

const (
	headerLength   = "Content-Length"
	headerType     = "Content-Type"
	headerEncoding = "Content-Encoding"
	headerVary     = "Vary"
)

type cryptoWriter struct {
	gin.ResponseWriter
	buffer *bytebufferpool.ByteBuffer
}

func (w *cryptoWriter) Write(b []byte) (int, error) {
	w.Header().Del(headerLength)
	w.Header().Del(headerType)
	w.Header().Del(headerEncoding)
	w.Header().Del(headerVary)
	return w.buffer.Write(b)
}

func (w *cryptoWriter) WriteString(s string) (int, error) {
	w.Header().Del(headerLength)
	w.Header().Del(headerType)
	w.Header().Del(headerEncoding)
	w.Header().Del(headerVary)
	return w.buffer.WriteString(s)
}

func (w *cryptoWriter) WriteHeader(code int) {
	w.Header().Del(headerLength)
	w.Header().Del(headerType)
	w.Header().Del(headerEncoding)
	w.Header().Del(headerVary)
	w.ResponseWriter.WriteHeader(code)
}

func Aes256Cbc(secret string) gin.HandlerFunc {
	key := []byte(cryptor.Md5String(secret))
	return func(c *gin.Context) {
		defer func(url string) {
			if err := recover(); err != nil {
				buf := make([]byte, 8192)
				n := runtime.Stack(buf, false)
				log.Error().Str("url", url).Interface("info", err).Str("stack", string(buf[:n])).Msg("Aes256Cbc panic")
				c.AbortWithStatus(http.StatusExpectationFailed)
			}
		}(c.Request.URL.String())

		buf := bytebufferpool.Get()
		defer bytebufferpool.Put(buf)

		rw := &cryptoWriter{buffer: bytebufferpool.Get(), ResponseWriter: c.Writer}
		defer bytebufferpool.Put(rw.buffer)

		n, err := buf.ReadFrom(c.Request.Body)
		if err == nil && n > 0 {
			if n < aes.BlockSize {
				c.AbortWithStatusJSON(http.StatusOK, gin.H{
					"success": false,
					"code":    90500,
					"message": "data length less than BlockSize",
				})
				return
			}
			_ = c.Request.Body.Close()
			data := cryptor.AesCbcDecrypt(buf.Bytes(), key)
			c.Request.Body = io.NopCloser(bytes.NewBuffer(data))
		}

		c.Request.Header.Set("Accept-Encoding", "gzip")
		c.Request.Header.Set("Content-Encoding", "gzip")
		c.Writer = rw

		c.Next()

		body := cryptor.AesCbcEncrypt(rw.buffer.Bytes(), key)
		c.Header(headerLength, fmt.Sprint(len(body)))
		c.Header(headerType, "application/octet-stream")

		_, err = rw.ResponseWriter.Write(body)
		if err != nil {
			log.Err(err).Msg("resp writer error")
		}
	}
}
